#include <iostream>
using namespace std;

struct TreeLinkNode {
    int val;
    TreeLinkNode *right;
    TreeLinkNode *left;
    TreeLinkNode *next;
    TreeLinkNode(int x) : val(x), left(NULL), right(NULL), next(NULL) {}
};

class Solution {
public:
    void connect(TreeLinkNode *root) {
        TreeLinkNode *p, *q;
        while(root != NULL) {
            p = root;
            root = NULL;

            // find the first non-empty child of root and replace root with it
            while(p != NULL) {
                if (p->left != NULL) {
                    root = p->left;
                    break;
                }
                if (p->right != NULL) {
                    root = p->right;
                    break;
                }
                p = p->next;
            }

            q = root;// After this, q becomes a child of p
            while(q != NULL) {
                TreeLinkNode *next = NULL;
                if (q == p->left && p->right != NULL) {
                    next = p->right;
                } else {
                    p = p->next;
                    while(p != NULL) {
                        if (p->left != NULL) {
                            next = p->left;
                            break;
                        }
                        if (p->right != NULL) {
                            next = p->right;
                            break;
                        }
                        p = p->next;
                    }
                }
                q->next = next;
                q = next;
            }
        }
    }
};
